import torch
from torch import nn
import torch.nn.functional as F


class CausalConv1d(nn.Module):
    """
    Input and output sizes will be the same, if stride == 1.
    $$
    y.shape[-1] = \frac{x.shape[-1] + kernel_size -2}{stride} - (kernel_size-1) + 1
    $$
    or
    (x.shape[-1] + kernel_size -2) // stride - (kernel_size-1) + 1
    x=101
    k, s = 6, 5
    print(x//s)
    print((x+k)//s - (k+x%s)//s)
    """

    def __init__(self, in_size, out_size, kernel_size, stride=1, dilation=1):
        super(CausalConv1d, self).__init__()
        self.pad = (kernel_size - 1) * dilation
        self.conv1 = nn.Conv1d(in_size, out_size, kernel_size, stride, padding=self.pad, dilation=dilation)

    def forward(self, x):
        x = self.conv1(x)
        x = x[..., :-self.pad]
        return x


class ResidualLayer(nn.Module):
    def __init__(self, residual_size, skip_size, dilation, kernel_size=3):
        super(ResidualLayer, self).__init__()
        self.conv_filter = CausalConv1d(residual_size, residual_size,
                                        kernel_size=kernel_size, dilation=dilation)
        self.conv_gate = CausalConv1d(residual_size, residual_size,
                                      kernel_size=kernel_size, dilation=dilation)
        self.resconv1_1 = nn.Conv1d(residual_size, residual_size, kernel_size=1)
        self.skipconv1_1 = nn.Conv1d(residual_size, skip_size, kernel_size=1)

    def forward(self, x):
        conv_filter = self.conv_filter(x)
        conv_gate = self.conv_gate(x)
        fx = torch.tanh(conv_filter) * torch.sigmoid(conv_gate)
        fx = self.resconv1_1(fx)
        skip = self.skipconv1_1(fx)
        residual = fx + x
        # residual=[batch,residual_size,seq_len]  skip=[batch,skip_size,seq_len]
        return skip, residual


class DilatedStack(nn.Module):
    def __init__(self, residual_size, skip_size, dilation_depth, kernel_size):
        super(DilatedStack, self).__init__()
        residual_stack = [ResidualLayer(residual_size, skip_size, 2 ** layer, kernel_size)
                          for layer in range(dilation_depth)]
        self.residual_stack = nn.ModuleList(residual_stack)

    def forward(self, x):
        skips = []
        for layer in self.residual_stack:
            skip, x = layer(x)
            skips.append(skip.unsqueeze(0))
            # skip =[1,batch,skip_size,seq_len]
        return torch.cat(skips, dim=0), x  # [layers,batch,skip_size,seq_len]


class WaveNet(nn.Module):

    def __init__(self, input_size, out_size, residual_size, skip_size, causal_stride=1, final_stride=1,
                 causal_kernel_size=3, dilated_kernel_size=3, dilation_cycles=3, dilation_depth=3):
        super(WaveNet, self).__init__()

        self.input_conv = CausalConv1d(input_size, residual_size, kernel_size=causal_kernel_size, stride=causal_stride)
        self.dilated_stacks = nn.ModuleList(

            [DilatedStack(residual_size, skip_size, dilation_depth, dilated_kernel_size)
             for cycle in range(dilation_cycles)]

        )

        self.convout_1 = nn.Conv1d(skip_size, skip_size, kernel_size=1)
        self.convout_2 = nn.Conv1d(skip_size, out_size, kernel_size=1, stride=final_stride)

    def forward(self, x):
        x = x.permute(0, 2, 1)  # [batch,input_feature_dim, seq_len]
        x = self.input_conv(x)  # [batch,residual_size, seq_len]

        skip_connections = []

        for cycle in self.dilated_stacks:
            skips, x = cycle(x)
            skip_connections.append(skips)

        ## skip_connection=[total_layers,batch,skip_size,seq_len]
        skip_connections = torch.cat(skip_connections, dim=0)

        # gather all output skip connections to generate output, discard last residual output

        out = skip_connections.sum(dim=0)  # [batch,skip_size,seq_len]

        out = F.relu(out)

        out = self.convout_1(out)  # [batch,skip_size,seq_len]
        out = F.relu(out)

        out = self.convout_2(out)

        out = out.permute(0, 2, 1)
        # [bacth,seq_len,out_size]
        return out